#!/usr/bin/env python3
# -*- coding:utf-8 -*-
# @Author: renjin@bit.edu.cn
# @Date  : 2024-08-06

"""
【节点名称】：
    COCODatasetLoaderNode
【依赖项安装】：
    pip install spirems
【订阅类型】：
    std_msgs::Boolean （是否输出下一张图像）
    spirecv_msgs::EvaluationResult （算法评价结果）
【发布类型】：
    sensor_msgs::CompressedImage （输出图像）
    std_msgs::Boolean （评价任务结束时，输出关闭所有节点指令）
【构造参数说明】：
    parameter_file (str): 全局参数文件
    sms_shutdown_emit (bool): 评价任务结束时，是否输出关闭所有节点指令
【节点参数】：
    coco_root_dir (str): 数据集根目录
    coco_gt_file (str): 数据集COCO类型真值
    img_total (int): 数据集图像数量
【备注】：
    无
"""

import threading
import time
import os
import json
from collections import defaultdict
import cv2
from typing import Union
from queue import Queue
from spirems import Subscriber, Publisher, cvimg2sms, def_msg, QoS
from spirecv.base.BaseNode import BaseNode
import uuid


class COCODatasetLoaderNode(threading.Thread, BaseNode):
    def __init__(
        self,
        job_name: str,
        ip: str = '127.0.0.1',
        port: int = 9094,
        param_dict_or_file: Union[dict, str] = None,
        sms_shutdown_emit: bool = True,
        sms_shutdown: bool = True,
        remote_ip: str = '127.0.0.1',
        remote_port: int = 9094
    ):
        threading.Thread.__init__(self)
        BaseNode.__init__(
            self,
            self.__class__.__name__,
            job_name,
            ip=ip,
            port=port,
            param_dict_or_file=param_dict_or_file,
            sms_shutdown=sms_shutdown
        )
        self.client_id = str(uuid.uuid4()).replace('-', '_')
        self.img_i_queue = Queue()
        self.queue_pool.append(self.img_i_queue)
        self.img_i = 0
        self.sms_shutdown_emit = sms_shutdown_emit

        self._next_reader = Subscriber(
            '/' + job_name + '/launch_next', 'std_msgs::Boolean', self.launch_next,
            ip=ip, port=port, qos=QoS.Reliability
        )

        evaluate_url = '/SpireView/{}/TaskResults'.format(self.client_id)
        self._evaluate_reader = Subscriber(
            evaluate_url, 'spirecv_msgs::EvaluationResult', self.evaluate_callback,
            ip=remote_ip, port=remote_port, qos=QoS.Reliability
        )
        self._evaluate_writer = Publisher(
            '/SpireView/EvaluationResult', 'spirecv_msgs::EvaluationResult',
            ip=ip, port=port, qos=QoS.Reliability
        )
        self._image_writer = Publisher(
            '/' + job_name + '/sensor/image_raw', 'sensor_msgs::CompressedImage',
            ip=ip, port=port, qos=QoS.Reliability
        )
        self._shutdown_writer = Publisher(
            '/' + job_name + '/shutdown', 'std_msgs::Boolean',
            ip=ip, port=port, qos=QoS.Reliability
        )

        self.coco_root_dir = self.get_param("coco_root_dir", r"D:\dataset\val2017")
        self.coco_gt_file = self.get_param("coco_gt_file", r"D:\dataset\annotations\instances_val2017.json")
        self.img_total = self.get_param("img_total", 5000)
        self._load_dataset()
        self.start()

    def release(self):
        BaseNode.release(self)
        self._next_reader.kill()
        self._evaluate_reader.kill()
        self._evaluate_writer.kill()
        self._image_writer.kill()
        self._shutdown_writer.kill()
        self._next_reader.join()
        self._evaluate_reader.join()
        self._evaluate_writer.join()
        self._image_writer.join()
        self._shutdown_writer.join()

    def launch_next(self, msg: dict = None):
        if (isinstance(msg, dict) and msg['data']) or msg is None:
            if self.img_i < len(self.img_keys):
                self.img_i_queue.put(self.img_i)
                self.img_i += 1

    def evaluate_callback(self, msg):
        print(msg)
        self._evaluate_writer.publish(msg)
        if self.sms_shutdown_emit:
            msg = def_msg('std_msgs::Boolean')
            msg['data'] = True
            self._shutdown_writer.publish(msg)

    def run(self):
        while self.is_running():
            img_i = self.img_i_queue.get(block=True)
            if img_i is None:
                break

            img_info = self.imgs[self.img_keys[img_i]]
            # img_id = img_info['id']
            file_name = os.path.basename(img_info['file_name'])
            img = cv2.imread(str(os.path.join(self.coco_root_dir, file_name)))

            assert img.shape[0] == img_info['height'] and img.shape[1] == img_info['width']
            # print(img.shape)
            msg = cvimg2sms(img, 'png')

            msg['img_id'] = img_i
            msg['img_total'] = self.img_total
            msg['file_name'] = file_name
            msg['client_id'] = self.client_id
            self._image_writer.publish(msg)

        self.release()
        print('{} quit!'.format(self.__class__.__name__))

    def _load_dataset(self):
        annotation_file = self.coco_gt_file
        tic = time.time()
        with open(annotation_file, 'r') as f:
            dataset = json.load(f)
        assert isinstance(dataset, dict), 'annotation file format {} not supported'.format(type(dataset))
        print('Done (t={:0.2f}s)'.format(time.time() - tic))
        self.dataset = dataset
        self._create_index()
        self.img_keys = list(self.imgs.keys())

    def _create_index(self):
        # create index
        print('creating index...')
        anns, cats, imgs = {}, {}, {}
        imgToAnns, catToImgs = defaultdict(list), defaultdict(list)
        if 'annotations' in self.dataset:
            for ann in self.dataset['annotations']:
                imgToAnns[ann['image_id']].append(ann)
                anns[ann['id']] = ann

        if 'images' in self.dataset:
            for img in self.dataset['images']:
                imgs[img['id']] = img

        if 'categories' in self.dataset:
            for cat in self.dataset['categories']:
                cats[cat['id']] = cat

        if 'annotations' in self.dataset and 'categories' in self.dataset:
            for ann in self.dataset['annotations']:
                catToImgs[ann['category_id']].append(ann['image_id'])

        print('index created!')
        # create class members
        self.anns = anns
        self.imgToAnns = imgToAnns
        self.catToImgs = catToImgs
        self.imgs = imgs
        self.cats = cats


if __name__ == '__main__':
    loader = COCODatasetLoaderNode("EvalJob")
